package com.fangdd.logtrace.http;

import com.fangdd.logtrace.BoundsLog;
import com.fangdd.logtrace.TraceId;
import com.fangdd.logtrace.TraceIdThreadLocal;
import org.apache.http.HttpHost;
import org.apache.http.HttpRequest;
import org.apache.http.HttpResponse;
import org.apache.http.client.ClientProtocolException;
import org.apache.http.client.HttpClient;
import org.apache.http.client.ResponseHandler;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.conn.ClientConnectionManager;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.params.HttpParams;
import org.apache.http.protocol.HttpContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;

/**
 * Created by sulong on 16-9-9.
 */
public class TraceIdHttpClient implements HttpClient, Closeable {
    private static final Logger LOGGER = LoggerFactory.getLogger(TraceIdHttpClient.class);

    private CloseableHttpClient closeableHttpClient;

    public TraceIdHttpClient(CloseableHttpClient closeableHttpClient) {
        if (closeableHttpClient == null) {
            throw new IllegalArgumentException(" closeableHttpClient is null ");
        }

        this.closeableHttpClient = closeableHttpClient;
    }

    @Override
    public void close() throws IOException {
        this.closeableHttpClient.close();
    }

    @Override
    public HttpParams getParams() {
        return this.closeableHttpClient.getParams();
    }

    @Override
    public ClientConnectionManager getConnectionManager() {
        return this.closeableHttpClient.getConnectionManager();
    }

    @Override
    public HttpResponse execute(final HttpUriRequest request) throws IOException, ClientProtocolException {
        return this.call(request, new HttpCall<HttpResponse>() {
            @Override
            public HttpResponse call() throws IOException {
                return closeableHttpClient.execute(request);
            }
        });
    }

    @Override
    public HttpResponse execute(final HttpUriRequest request, final HttpContext context) throws IOException, ClientProtocolException {
        return this.call(request, new HttpCall<HttpResponse>() {
            @Override
            public HttpResponse call() throws IOException {
                return closeableHttpClient.execute(request, context);
            }
        });
    }

    @Override
    public HttpResponse execute(final HttpHost target, final HttpRequest request) throws IOException, ClientProtocolException {
        return this.call(request, new HttpCall<HttpResponse>() {
            @Override
            public HttpResponse call() throws IOException {
                return closeableHttpClient.execute(target, request);
            }
        });
    }

    @Override
    public HttpResponse execute(final HttpHost target, final HttpRequest request, final HttpContext context) throws IOException, ClientProtocolException {
        return this.call(request, new HttpCall<HttpResponse>() {
            @Override
            public HttpResponse call() throws IOException {
                return closeableHttpClient.execute(target, request, context);
            }
        });
    }

    private <R> R call(HttpRequest request ,  HttpCall<R> call) throws IOException {
        TraceId next = TraceIdThreadLocal.sequentialNext();
        String method = "http:" + request.getRequestLine().getMethod() + ":" + request.getRequestLine().getUri();
        if (!next.isUnset()) {
            request.setHeader(TraceIdHttpFilter.TRACE_ID_HEADER_NAME, next.toString());
            BoundsLog.callerIn(LOGGER, next.toString(), method);
        }

        long startTime = System.currentTimeMillis();
        try {
            R httpResponse = call.call();
            if (!next.isUnset()) {
                BoundsLog.callerOut(LOGGER, next.toString(), method, true, System.currentTimeMillis() - startTime);
            }
            return httpResponse;
        } catch (ClientProtocolException e) {
            if (!next.isUnset()) {
                BoundsLog.callerOut(LOGGER, next.toString(), method, true, System.currentTimeMillis() - startTime);
            }
            throw e;
        } catch (IOException ioException) {
            if (!next.isUnset()) {
                BoundsLog.callerOut(LOGGER, next.toString(), method, false, System.currentTimeMillis() - startTime);
            }
            throw ioException;
        }
    }

    @Override
    public <T> T execute(final HttpUriRequest request, final ResponseHandler<? extends T> responseHandler) throws IOException, ClientProtocolException {
        return this.call(request, new HttpCall<T>() {
            @Override
            public T call() throws IOException {
                return closeableHttpClient.execute(request, responseHandler);
            }
        });
    }

    @Override
    public <T> T execute(final HttpUriRequest request, final ResponseHandler<? extends T> responseHandler, final HttpContext context) throws IOException, ClientProtocolException {
        return this.call(request, new HttpCall<T>() {
            @Override
            public T call() throws IOException {
                return closeableHttpClient.execute(request, responseHandler, context);
            }
        });
    }

    @Override
    public <T> T execute(final HttpHost target, final HttpRequest request, final ResponseHandler<? extends T> responseHandler) throws IOException, ClientProtocolException {
        return this.call(request, new HttpCall<T>() {
            @Override
            public T call() throws IOException {
                return closeableHttpClient.execute(target, request, responseHandler);
            }
        });
    }

    @Override
    public <T> T execute(final HttpHost target, final HttpRequest request, final ResponseHandler<? extends T> responseHandler, final HttpContext context) throws IOException, ClientProtocolException {
        return this.call(request, new HttpCall<T>() {
            @Override
            public T call() throws IOException {
                return closeableHttpClient.execute(target, request, responseHandler, context);
            }
        });
    }

    private interface HttpCall<R> {
        R call() throws IOException;
    }
}
